from GStreamerWrapper import GStreamerWrapper
import cv2
import os
from datetime import datetime, timezone


def ensure_output_dir() -> str:
    current_dir = os.path.dirname(os.path.abspath(__file__))
    output_dir = os.path.join(current_dir, "test")
    os.makedirs(output_dir, exist_ok=True)
    return output_dir


def generate_utc_filename() -> str:
    # 以世界时间(UTC)命名，格式：YYYYMMDDTHHMMSSZ.jpg
    return datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ") + ".jpg"


def main() -> None:
    print("Initializing GStreamerWrapper with 1920x1080...")
    wrapper = GStreamerWrapper(width=1920, height=1080)

    # 等待第一帧
    print("Waiting for the first frame...")
    frame = None
    while frame is None:
        frame = wrapper.GetFrame()

    output_dir = ensure_output_dir()
    print("\n------------------------------")
    print("准备就绪。按 回车 保存图片；输入 q 后按 回车 退出。")
    print(f"保存目录: {output_dir}")
    print("------------------------------")

    try:
        while True:
            user_input = input("")
            if user_input.strip().lower() == "q":
                break

            frame = wrapper.GetFrame()
            if frame is None:
                print("未获取到图像帧，稍后重试...")
                continue

            filename = generate_utc_filename()
            save_path = os.path.join(output_dir, filename)
            success = cv2.imwrite(save_path, frame)
            if success:
                print(f"已保存: {save_path}")
            else:
                print("保存失败，请重试。")
    except KeyboardInterrupt:
        pass
    finally:
        print("Stopping GStreamer thread...")
        wrapper.StopThread()
        print("Done.")


if __name__ == "__main__":
    main()


